#!/usr/bin/env python3

import udp
import sys
import packet
import os
import getpass
import mmap

def parseArgs(argv):
    argc = len(argv)
    if argc < 2:
        return print("请指定服务器及验证信息!")

    data = {}
    mes = argv[1]

    #  解析用户名
    if '@' in mes:
        data['user'], mes = mes.split('@')
    else:
        data['user'] = os.getlogin()

    #  参数中没有:代表参数格式错误
    if ':' not in mes:
        return print("请指定要下载的文件路径")

    #  解析IP port和下载文件路径
    ip, *mes = mes.split(':')
    if len(mes) == 2:
        port, data['path'] = mes
    else:
        data['path'] = mes[0]
        port = 9000

    #  判断下载文件路径是否为绝对路径
    if data['path'][0] != '/':
        return print("请正确指定要下载的文件路径(绝对路径)")

    #  规范化下载文件路径
    data['path'] = os.path.realpath(data['path'])

    #  解析目标存储文件路径
    if argc >= 3:
        dstpath = argv[2]
        if os.path.isdir(dstpath):
            dstpath += '/' + os.path.basename(data['path'])
            dstpath = os.path.realpath(dstpath)
    else:
        dstpath = os.path.basename(data['path'])

    if os.path.exists(dstpath) and input("是否覆盖(Y/N): ") not in "Yy":
        return None

    data['pwd'] = getpass.getpass()
    if not data['pwd']:
        return print("请好好输入密码...")
    #  data['pwd'] = '1'

    return ip, int(port), dstpath, data

def exitStr(sd, s):
    sd.close()
    return print(s)

def main():
    data = parseArgs(sys.argv)
    if not data:
        return
    ip, port, dstpath, data = data

    #  创建套接字和发送连接数据包
    sd = udp.cli_create((ip, port), data)
    if not sd:
        return print("连接服务器失败!")

    #  接受连接是否成功包
    data = packet.recv(sd, timeout=5.0)
    if not data:
        return exitStr(sd, "连接服务器失败!")

    data, addr = data
    if data['type'] == packet.ERROR:
        return exitStr(sd, packet.errorMessage[data['code']])

    #  打开本地存储的文件
    f = open(dstpath, "wb+")
    #  f.truncate(data['size'])
    #  f.flush()
    #  m = mmap.mmap(f.fileno(), data['size'])

    count = 0
    while True:
        #  接受数据包
        data = packet.recv(sd, timeout=5.0)
        if not data:
            return exitStr(sd, "接受数据包失败")
        data, addr = data
        if data['type'] == packet.DATA:
            #  print(count, len(data['data']))
            #  将数据写入文件　
            f.write(data['data'])
            #  给服务器返回数据回应包
            if not packet.send(sd, None, addr, packet.DATA_ACK):
                return exitStr(sd, "发送数据回应包失败")
            count += 1
            if len(data['data']) != 512:
                break
        elif data['type'] == packet.SHUTDOWN:
            #  判断退出条件退出循环
            break

    #  m.flush()
    #  m.close()
    f.flush()
    f.close()
    sd.close()

if __name__ == "__main__":
    main()
